# -*- coding: utf-8 -*-
"""
@Time ： 2024/3/5 14:50
@Auth ： gc
"""
import baostock as bs
import pandas as pd
import matplotlib.pyplot as plt
import datetime
import matplotlib.dates as mdates
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False    # 用来正常显示负号


def z_score_normalize(data_series):
    mean = data_series.mean()
    std = data_series.std()
    normalized = (data_series - mean) / std
    return normalized

# 复权类型，默认不复权：3；1：后复权；2：前复权。
def get_hs300_zs_info(start_date, end_date, code = "sh.000300", adjustflag = "3"):
    """
    获取沪深300指数情况
    :return:
    """
    # 登陆系统
    lg = bs.login()
    rs = bs.query_history_k_data_plus(code,
                                      "date,code,open,high,low,close,preclose,volume,amount,pctChg",
                                      start_date = start_date, end_date = end_date, frequency = "d", adjustflag = adjustflag)
    data_list = []
    while (rs.error_code == '0') & rs.next():
        # 获取一条记录，将记录合并在一起
        data_list.append(rs.get_row_data())
    result = pd.DataFrame(data_list, columns = rs.fields)
    result["close"] = result["close"].astype(float)
    result["date"] = result["date"].astype("datetime64[ns]")
    return result


def comprehensive_analysis(code = "sh.000300", adjustflag = "3"):
    now = datetime.datetime.now()
    pre_month = now - datetime.timedelta(days = 30)
    pre_60 = now - datetime.timedelta(days = 60)
    pre_year = now - datetime.timedelta(days = 365)
    pre_3_year = now - datetime.timedelta(days = 365 * 3)

    now_str = now.strftime("%Y-%m-%d")
    pre_month_str = pre_month.strftime("%Y-%m-%d")
    pre_60_str = pre_60.strftime("%Y-%m-%d")
    start_year_str = str(now.year) + "-01-01"
    pre_year_str = pre_year.strftime("%Y-%m-%d")
    pre_3_year_str = pre_3_year.strftime("%Y-%m-%d")
    
    
    datas_2015 = get_hs300_zs_info(start_date = '2015-01-01', end_date = now_str, code = code, adjustflag = adjustflag)
    
    
    datas = get_hs300_zs_info(start_date = pre_3_year_str, end_date = now_str, code = code, adjustflag = adjustflag)

    fig, axes = plt.subplots(6, 1, figsize = (60, 40))
    

    axes[5].plot(datas_2015["date"], datas_2015["close"])
    axes[5].plot(datas_2015["date"], [round(datas_2015["close"].describe()["25%"], 2)]*datas_2015.shape[0], linestyle = '--', color = 'gray')
    axes[5].plot(datas_2015["date"], [round(datas_2015["close"].describe()["50%"], 2)]*datas_2015.shape[0], linestyle = '--', color = 'green')
    axes[5].plot(datas_2015["date"], [round(datas_2015["close"].describe()["mean"], 2)]*datas_2015.shape[0], linestyle = '--', color = 'red')
    axes[5].plot(datas_2015["date"], [round(datas_2015["close"].describe()["75%"], 2)]*datas_2015.shape[0], linestyle = '--', color = 'purple')
    axes[5].text(datas_2015["date"].iloc[-1], round(datas_2015["close"].describe()["25%"], 2), f'25%_{round(datas_2015["close"].describe()["25%"], 2)}')
    axes[5].text(datas_2015["date"].iloc[-1], round(datas_2015["close"].describe()["50%"], 2), f'50%_{round(datas_2015["close"].describe()["50%"], 2)}')
    axes[5].text(datas_2015["date"].iloc[-1], round(datas_2015["close"].describe()["mean"], 2), f'mean_{round(datas_2015["close"].describe()["mean"], 2)}')
    axes[5].text(datas_2015["date"].iloc[-1], round(datas_2015["close"].describe()["75%"], 2), f'75%_{round(datas_2015["close"].describe()["75%"], 2)}')
    axes[5].xaxis.set_major_locator(mdates.YearLocator())
    axes[5].xaxis.set_major_formatter(mdates.DateFormatter('%Y'))
    axes[5].set_title("15年走势")
    axes[5].legend()
    
    
    axes[0].plot(datas["date"], datas["close"])
    axes[0].plot(datas["date"], [round(datas["close"].describe()["25%"], 2)]*datas.shape[0], linestyle = '--', color = 'gray')
    axes[0].plot(datas["date"], [round(datas["close"].describe()["50%"], 2)]*datas.shape[0], linestyle = '--', color = 'green')
    axes[0].plot(datas["date"], [round(datas["close"].describe()["mean"], 2)]*datas.shape[0], linestyle = '--', color = 'red')
    axes[0].plot(datas["date"], [round(datas["close"].describe()["75%"], 2)]*datas.shape[0], linestyle = '--', color = 'purple')
    axes[0].text(datas["date"].iloc[-1], round(datas["close"].describe()["25%"], 2), f'25%_{round(datas["close"].describe()["25%"], 2)}')
    axes[0].text(datas["date"].iloc[-1], round(datas["close"].describe()["50%"], 2), f'50%_{round(datas["close"].describe()["50%"], 2)}')
    axes[0].text(datas["date"].iloc[-1], round(datas["close"].describe()["mean"], 2), f'mean_{round(datas["close"].describe()["mean"], 2)}')
    axes[0].text(datas["date"].iloc[-1], round(datas["close"].describe()["75%"], 2), f'75%_{round(datas["close"].describe()["75%"], 2)}')
    axes[0].xaxis.set_major_locator(mdates.MonthLocator())
    axes[0].xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m'))
    axes[0].set_title("三年走势")
    axes[0].legend()

    j = datas[datas["date"] >= pre_year_str]
    axes[1].plot(j["date"], j["close"])
    axes[1].plot(j["date"], [round(j["close"].describe()["25%"], 2)]*j.shape[0], linestyle = '--', color = 'gray')
    axes[1].plot(j["date"], [round(j["close"].describe()["50%"], 2)]*j.shape[0], linestyle = '--', color = 'green')
    axes[1].plot(j["date"], [round(j["close"].describe()["mean"], 2)]*j.shape[0], linestyle = '--', color = 'red')
    axes[1].plot(j["date"], [round(j["close"].describe()["75%"], 2)]*j.shape[0], linestyle = '--', color = 'purple')
    axes[1].text(j["date"].iloc[-1], round(j["close"].describe()["25%"], 2), f'25%_{round(j["close"].describe()["25%"], 2)}')
    axes[1].text(j["date"].iloc[-1], round(j["close"].describe()["50%"], 2), f'50%_{round(j["close"].describe()["50%"], 2)}')
    axes[1].text(j["date"].iloc[-1], round(j["close"].describe()["mean"], 2), f'mean_{round(j["close"].describe()["mean"], 2)}')
    axes[1].text(j["date"].iloc[-1], round(j["close"].describe()["75%"], 2), f'75%_{round(j["close"].describe()["75%"], 2)}')
    axes[1].xaxis.set_major_locator(mdates.MonthLocator())
    axes[1].xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m'))
    axes[1].set_title("一年走势")
    axes[1].legend()

    j = datas[datas["date"] >= start_year_str]
    axes[2].plot(j["date"], j["close"])
    axes[2].plot(j["date"], [round(j["close"].describe()["25%"], 2)]*j.shape[0], linestyle = '--', color = 'gray')
    axes[2].plot(j["date"], [round(j["close"].describe()["50%"], 2)]*j.shape[0], linestyle = '--', color = 'green')
    axes[2].plot(j["date"], [round(j["close"].describe()["mean"], 2)]*j.shape[0], linestyle = '--', color = 'red')
    axes[2].plot(j["date"], [round(j["close"].describe()["75%"], 2)]*j.shape[0], linestyle = '--', color = 'purple')
    axes[2].text(j["date"].iloc[-1], round(j["close"].describe()["25%"], 2), f'25%_{round(j["close"].describe()["25%"], 2)}')
    axes[2].text(j["date"].iloc[-1], round(j["close"].describe()["50%"], 2), f'50%_{round(j["close"].describe()["50%"], 2)}')
    axes[2].text(j["date"].iloc[-1], round(j["close"].describe()["mean"], 2), f'mean_{round(j["close"].describe()["mean"], 2)}')
    axes[2].text(j["date"].iloc[-1], round(j["close"].describe()["75%"], 2), f'75%_{round(j["close"].describe()["75%"], 2)}')
    axes[2].xaxis.set_major_locator(mdates.MonthLocator())
    axes[2].xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m'))
    axes[2].set_title("年初到现在走势")
    axes[2].legend()

    j = datas[datas["date"] >= pre_60_str]
    axes[3].plot(j["date"], j["close"])
    axes[3].plot(j["date"], [round(j["close"].describe()["25%"], 2)]*j.shape[0], linestyle = '--', color = 'gray')
    axes[3].plot(j["date"], [round(j["close"].describe()["50%"], 2)]*j.shape[0], linestyle = '--', color = 'green')
    axes[3].plot(j["date"], [round(j["close"].describe()["mean"], 2)]*j.shape[0], linestyle = '--', color = 'red')
    axes[3].plot(j["date"], [round(j["close"].describe()["75%"], 2)]*j.shape[0], linestyle = '--', color = 'purple')
    axes[3].text(j["date"].iloc[-1], round(j["close"].describe()["25%"], 2), f'25%_{round(j["close"].describe()["25%"], 2)}')
    axes[3].text(j["date"].iloc[-1], round(j["close"].describe()["50%"], 2), f'50%_{round(j["close"].describe()["50%"], 2)}')
    axes[3].text(j["date"].iloc[-1], round(j["close"].describe()["mean"], 2), f'mean_{round(j["close"].describe()["mean"], 2)}')
    axes[3].text(j["date"].iloc[-1], round(j["close"].describe()["75%"], 2), f'75%_{round(j["close"].describe()["75%"], 2)}')
    axes[3].xaxis.set_major_locator(mdates.WeekdayLocator())
    # axes[3].xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-d%'))
    axes[3].set_title("最近2月走势")
    axes[3].legend()

    j = datas[datas["date"] >= pre_month_str]
    axes[4].plot(j["date"], j["close"])
    axes[4].plot(j["date"], [round(j["close"].describe()["25%"], 2)] * j.shape[0], linestyle = '--', color = 'gray')
    axes[4].plot(j["date"], [round(j["close"].describe()["50%"], 2)] * j.shape[0], linestyle = '--', color = 'green')
    axes[4].plot(j["date"], [round(j["close"].describe()["mean"], 2)] * j.shape[0], linestyle = '--', color = 'red')
    axes[4].plot(j["date"], [round(j["close"].describe()["75%"], 2)] * j.shape[0], linestyle = '--', color = 'purple')
    axes[4].text(j["date"].iloc[-1], round(j["close"].describe()["25%"], 2), f'25%_{round(j["close"].describe()["25%"], 2)}')
    axes[4].text(j["date"].iloc[-1], round(j["close"].describe()["50%"], 2), f'50%_{round(j["close"].describe()["50%"], 2)}')
    axes[4].text(j["date"].iloc[-1], round(j["close"].describe()["mean"], 2), f'mean_{round(j["close"].describe()["mean"], 2)}')
    axes[4].text(j["date"].iloc[-1], round(j["close"].describe()["75%"], 2), f'75%_{round(j["close"].describe()["75%"], 2)}')
    axes[4].xaxis.set_major_locator(mdates.WeekdayLocator())
    # axes[4].xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-d%'))
    axes[4].set_title("最近1月走势")
    axes[4].legend()
    plt.savefig(f"fund/data/zs_{code.replace('.', '')}_info_{adjustflag}.png")
    # plt.legend()
    # plt.show()


def stock_comprehensive_analysis(code = "sh.000300", adjustflag = "3", use_15 = True):
    now = datetime.datetime.now()
    pre_month = now - datetime.timedelta(days = 30)
    pre_60 = now - datetime.timedelta(days = 60)
    pre_year = now - datetime.timedelta(days = 365)
    pre_3_year = now - datetime.timedelta(days = 365 * 3)

    now_str = now.strftime("%Y-%m-%d")
    pre_month_str = pre_month.strftime("%Y-%m-%d")
    pre_60_str = pre_60.strftime("%Y-%m-%d")
    start_year_str = str(now.year) + "-01-01"
    pre_year_str = pre_year.strftime("%Y-%m-%d")
    pre_3_year_str = pre_3_year.strftime("%Y-%m-%d")
    if use_15:
        datas_2015 = get_hs300_zs_info(start_date = '2015-01-01', end_date = now_str, code = "sh.000300", adjustflag = adjustflag)
        stock_datas_2015 = get_hs300_zs_info(start_date = '2015-01-01', end_date = now_str, code = code, adjustflag = adjustflag)
        df_2015 = stock_datas_2015["close"]/datas_2015["close"]
        datas_2015["close"] = (df_2015-df_2015.mean())/df_2015.std()

    datas = get_hs300_zs_info(start_date = pre_3_year_str, end_date = now_str, code = "sh.000300", adjustflag = adjustflag)
    stock_datas = get_hs300_zs_info(start_date = pre_3_year_str, end_date = now_str, code = code, adjustflag = adjustflag)
    df = stock_datas["close"] / datas["close"]
    datas["close"] = (df - df.mean()) / df.std()

    fig, axes = plt.subplots(6, 1, figsize = (60, 40))
    if use_15:
        axes[5].plot(datas_2015["date"], datas_2015["close"])
        axes[5].plot(datas_2015["date"], [round(datas_2015["close"].describe()["25%"], 2)] * datas_2015.shape[0],
                     linestyle = '--', color = 'gray')
        axes[5].plot(datas_2015["date"], [round(datas_2015["close"].describe()["50%"], 2)] * datas_2015.shape[0],
                     linestyle = '--', color = 'green')
        axes[5].plot(datas_2015["date"], [round(datas_2015["close"].describe()["mean"], 2)] * datas_2015.shape[0],
                     linestyle = '--', color = 'red')
        axes[5].plot(datas_2015["date"], [round(datas_2015["close"].describe()["75%"], 2)] * datas_2015.shape[0],
                     linestyle = '--', color = 'purple')
        axes[5].text(datas_2015["date"].iloc[-1], round(datas_2015["close"].describe()["25%"], 2),
                     f'25%_{round(datas_2015["close"].describe()["25%"], 2)}')
        axes[5].text(datas_2015["date"].iloc[-1], round(datas_2015["close"].describe()["50%"], 2),
                     f'50%_{round(datas_2015["close"].describe()["50%"], 2)}')
        axes[5].text(datas_2015["date"].iloc[-1], round(datas_2015["close"].describe()["mean"], 2),
                     f'mean_{round(datas_2015["close"].describe()["mean"], 2)}')
        axes[5].text(datas_2015["date"].iloc[-1], round(datas_2015["close"].describe()["75%"], 2),
                     f'75%_{round(datas_2015["close"].describe()["75%"], 2)}')
        axes[5].xaxis.set_major_locator(mdates.YearLocator())
        axes[5].xaxis.set_major_formatter(mdates.DateFormatter('%Y'))
        axes[5].set_title("15年走势")
        axes[5].legend()

    axes[0].plot(datas["date"], datas["close"])
    axes[0].plot(datas["date"], [round(datas["close"].describe()["25%"], 2)] * datas.shape[0], linestyle = '--',
                 color = 'gray')
    axes[0].plot(datas["date"], [round(datas["close"].describe()["50%"], 2)] * datas.shape[0], linestyle = '--',
                 color = 'green')
    axes[0].plot(datas["date"], [round(datas["close"].describe()["mean"], 2)] * datas.shape[0], linestyle = '--',
                 color = 'red')
    axes[0].plot(datas["date"], [round(datas["close"].describe()["75%"], 2)] * datas.shape[0], linestyle = '--',
                 color = 'purple')
    axes[0].text(datas["date"].iloc[-1], round(datas["close"].describe()["25%"], 2),
                 f'25%_{round(datas["close"].describe()["25%"], 2)}')
    axes[0].text(datas["date"].iloc[-1], round(datas["close"].describe()["50%"], 2),
                 f'50%_{round(datas["close"].describe()["50%"], 2)}')
    axes[0].text(datas["date"].iloc[-1], round(datas["close"].describe()["mean"], 2),
                 f'mean_{round(datas["close"].describe()["mean"], 2)}')
    axes[0].text(datas["date"].iloc[-1], round(datas["close"].describe()["75%"], 2),
                 f'75%_{round(datas["close"].describe()["75%"], 2)}')
    axes[0].xaxis.set_major_locator(mdates.MonthLocator())
    axes[0].xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m'))
    axes[0].set_title("三年走势")
    axes[0].legend()

    j = datas[datas["date"] >= pre_year_str]
    axes[1].plot(j["date"], j["close"])
    axes[1].plot(j["date"], [round(j["close"].describe()["25%"], 2)] * j.shape[0], linestyle = '--', color = 'gray')
    axes[1].plot(j["date"], [round(j["close"].describe()["50%"], 2)] * j.shape[0], linestyle = '--', color = 'green')
    axes[1].plot(j["date"], [round(j["close"].describe()["mean"], 2)] * j.shape[0], linestyle = '--', color = 'red')
    axes[1].plot(j["date"], [round(j["close"].describe()["75%"], 2)] * j.shape[0], linestyle = '--', color = 'purple')
    axes[1].text(j["date"].iloc[-1], round(j["close"].describe()["25%"], 2),
                 f'25%_{round(datas["close"].describe()["25%"], 2)}')
    axes[1].text(j["date"].iloc[-1], round(j["close"].describe()["50%"], 2),
                 f'50%_{round(datas["close"].describe()["50%"], 2)}')
    axes[1].text(j["date"].iloc[-1], round(j["close"].describe()["mean"], 2),
                 f'mean_{round(datas["close"].describe()["mean"], 2)}')
    axes[1].text(j["date"].iloc[-1], round(j["close"].describe()["75%"], 2),
                 f'75%_{round(datas["close"].describe()["75%"], 2)}')
    axes[1].xaxis.set_major_locator(mdates.MonthLocator())
    axes[1].xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m'))
    axes[1].set_title("一年走势")
    axes[1].legend()

    j = datas[datas["date"] >= start_year_str]
    axes[2].plot(j["date"], j["close"])
    axes[2].plot(j["date"], [round(j["close"].describe()["25%"], 2)] * j.shape[0], linestyle = '--', color = 'gray')
    axes[2].plot(j["date"], [round(j["close"].describe()["50%"], 2)] * j.shape[0], linestyle = '--', color = 'green')
    axes[2].plot(j["date"], [round(j["close"].describe()["mean"], 2)] * j.shape[0], linestyle = '--', color = 'red')
    axes[2].plot(j["date"], [round(j["close"].describe()["75%"], 2)] * j.shape[0], linestyle = '--', color = 'purple')
    axes[2].text(j["date"].iloc[-1], round(j["close"].describe()["25%"], 2),
                 f'25%_{round(datas["close"].describe()["25%"], 2)}')
    axes[2].text(j["date"].iloc[-1], round(j["close"].describe()["50%"], 2),
                 f'50%_{round(datas["close"].describe()["50%"], 2)}')
    axes[2].text(j["date"].iloc[-1], round(j["close"].describe()["mean"], 2),
                 f'mean_{round(datas["close"].describe()["mean"], 2)}')
    axes[2].text(j["date"].iloc[-1], round(j["close"].describe()["75%"], 2),
                 f'75%_{round(datas["close"].describe()["75%"], 2)}')
    axes[2].xaxis.set_major_locator(mdates.MonthLocator())
    axes[2].xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m'))
    axes[2].set_title("年初到现在走势")
    axes[2].legend()

    j = datas[datas["date"] >= pre_60_str]
    axes[3].plot(j["date"], j["close"])
    axes[3].plot(j["date"], [round(j["close"].describe()["25%"], 2)] * j.shape[0], linestyle = '--', color = 'gray')
    axes[3].plot(j["date"], [round(j["close"].describe()["50%"], 2)] * j.shape[0], linestyle = '--', color = 'green')
    axes[3].plot(j["date"], [round(j["close"].describe()["mean"], 2)] * j.shape[0], linestyle = '--', color = 'red')
    axes[3].plot(j["date"], [round(j["close"].describe()["75%"], 2)] * j.shape[0], linestyle = '--', color = 'purple')
    axes[3].text(j["date"].iloc[-1], round(j["close"].describe()["25%"], 2),
                 f'25%_{round(datas["close"].describe()["25%"], 2)}')
    axes[3].text(j["date"].iloc[-1], round(j["close"].describe()["50%"], 2),
                 f'50%_{round(datas["close"].describe()["50%"], 2)}')
    axes[3].text(j["date"].iloc[-1], round(j["close"].describe()["mean"], 2),
                 f'mean_{round(datas["close"].describe()["mean"], 2)}')
    axes[3].text(j["date"].iloc[-1], round(j["close"].describe()["75%"], 2),
                 f'75%_{round(datas["close"].describe()["75%"], 2)}')
    axes[3].xaxis.set_major_locator(mdates.WeekdayLocator())
    # axes[3].xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-d%'))
    axes[3].set_title("最近2月走势")
    axes[3].legend()

    j = datas[datas["date"] >= pre_month_str]
    axes[4].plot(j["date"], j["close"])
    axes[4].plot(j["date"], [round(j["close"].describe()["25%"], 2)] * j.shape[0], linestyle = '--', color = 'gray')
    axes[4].plot(j["date"], [round(j["close"].describe()["50%"], 2)] * j.shape[0], linestyle = '--', color = 'green')
    axes[4].plot(j["date"], [round(j["close"].describe()["mean"], 2)] * j.shape[0], linestyle = '--', color = 'red')
    axes[4].plot(j["date"], [round(j["close"].describe()["75%"], 2)] * j.shape[0], linestyle = '--', color = 'purple')
    axes[4].text(j["date"].iloc[-1], round(j["close"].describe()["25%"], 2),
                 f'25%_{round(datas["close"].describe()["25%"], 2)}')
    axes[4].text(j["date"].iloc[-1], round(j["close"].describe()["50%"], 2),
                 f'50%_{round(datas["close"].describe()["50%"], 2)}')
    axes[4].text(j["date"].iloc[-1], round(j["close"].describe()["mean"], 2),
                 f'mean_{round(datas["close"].describe()["mean"], 2)}')
    axes[4].text(j["date"].iloc[-1], round(j["close"].describe()["75%"], 2),
                 f'75%_{round(datas["close"].describe()["75%"], 2)}')
    axes[4].xaxis.set_major_locator(mdates.WeekdayLocator())
    # axes[4].xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-d%'))
    axes[4].set_title("最近1月走势")
    axes[4].legend()
    plt.savefig(f"fund/data/stock_{code.replace('.', '')}_info_{adjustflag}.png")


if __name__ == '__main__':
    # get_hs300_zs_info("2024-01-01", "2024-03-19",adjustflag="2")
    comprehensive_analysis()